#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/time.h>

#define PACKET_SIZE 64
#define PORT 0
#define TIMEOUT_SEC 2

// 计算校验和
unsigned short checksum(void *b, int len) {    
    unsigned short *buf = (unsigned short *)b;
    unsigned int sum = 0;
    unsigned short result;

    for (sum = 0; len > 1; len -= 2)
        sum += *buf++;

    if (len == 1)
        sum += *(unsigned char *)buf;

    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    result = ~sum;

    return result;
}

// 发送 ICMP Echo 请求
void send_ping(int sockfd, struct sockaddr *addr) {
    char sendbuf[PACKET_SIZE];
    struct icmp *icmp_hdr;
    int packet_size;

    icmp_hdr = (struct icmp *)sendbuf;
    icmp_hdr->icmp_type = ICMP_ECHO;
    icmp_hdr->icmp_code = 0;
    icmp_hdr->icmp_id = getpid();
    icmp_hdr->icmp_seq = 0;
    memset(icmp_hdr->icmp_data, 0xa5, PACKET_SIZE - sizeof(struct icmp));
    gettimeofday((struct timeval *)icmp_hdr->icmp_data, NULL);

    packet_size = 8 + PACKET_SIZE;
    icmp_hdr->icmp_cksum = 0;
    icmp_hdr->icmp_cksum = checksum(icmp_hdr, packet_size);

    if (sendto(sockfd, sendbuf, packet_size, 0, addr, sizeof(*addr)) <= 0) {
        perror("sendto");
        exit(EXIT_FAILURE);
    }
}

// 接收 ICMP Echo 响应
void receive_ping(int sockfd, struct sockaddr *addr) {
    char recvbuf[PACKET_SIZE + sizeof(struct iphdr)];
    struct iphdr *ip_hdr;
    struct icmp *icmp_hdr;
    struct sockaddr_in from;
    socklen_t fromlen = sizeof(from);
    struct timeval end_time;

    gettimeofday(&end_time, NULL);
    end_time.tv_sec += TIMEOUT_SEC;

    while (1) {
        fd_set fdset;
        struct timeval timeout;
        int ret;

        FD_ZERO(&fdset);
        FD_SET(sockfd, &fdset);

        gettimeofday(&timeout, NULL);
        timeout.tv_sec = end_time.tv_sec - timeout.tv_sec;
        timeout.tv_usec = end_time.tv_usec - timeout.tv_usec;

        ret = select(sockfd + 1, &fdset, NULL, NULL, &timeout);
        if (ret == -1) {
            perror("select");
            exit(EXIT_FAILURE);
        } else if (ret == 0) {
            printf("Timeout\n");
            exit(EXIT_SUCCESS);
        }

        if (recvfrom(sockfd, recvbuf, sizeof(recvbuf), 0, (struct sockaddr *)&from, &fromlen) <= 0) {
            perror("recvfrom");
            exit(EXIT_FAILURE);
        }

        ip_hdr = (struct iphdr *)recvbuf;
        icmp_hdr = (struct icmp *)(recvbuf + (ip_hdr->ihl << 2));

        if (icmp_hdr->icmp_type == ICMP_ECHOREPLY) {
            printf("Reply from %s: icmp_seq=%d\n", inet_ntoa(from.sin_addr), icmp_hdr->icmp_seq);
            break;
        }
    }
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <hostname or IP address>\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    struct hostent *host;
    struct sockaddr_in addr;
    int sockfd;

    if ((sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)) < 0) {
        perror("socket");
        exit(EXIT_FAILURE);
    }

    if ((host = gethostbyname(argv[1])) == NULL) {
        perror("gethostbyname");
        exit(EXIT_FAILURE);
    }

    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = *((struct in_addr *)host->h_addr);

    printf("Pinging %s ...\n", inet_ntoa(addr.sin_addr));

    send_ping(sockfd, (struct sockaddr *)&addr);
    receive_ping(sockfd, (struct sockaddr *)&addr);

    close(sockfd);

    return 0;
}